import wmi
import mysql.connector
from datetime import datetime
import subprocess
import re
import platform

class NodeManager:
    def __init__(self, conn):
        self.conn = conn
        
    def get_disk_serial(self):
        """获取硬盘序列号"""
        try:
            if platform.system() == 'Windows':
                # Windows系统使用WMI获取
                c = wmi.WMI()
                for physical_disk in c.Win32_DiskDrive():
                    if physical_disk.SerialNumber:
                        return physical_disk.SerialNumber.strip()
            else:
                # Linux系统使用hdparm命令获取
                cmd = "sudo hdparm -I /dev/sda | grep 'Serial Number'"
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
                if result.stdout:
                    serial = re.search(r'Serial Number:\s*(.+)', result.stdout)
                    if serial:
                        return serial.group(1).strip()
            
            return "UNKNOWN"
        except Exception as e:
            print(f"Error getting disk serial: {e}")
            return "ERROR"

    def update_node_status(self):
        """更新节点状态"""
        try:
            # 获取硬盘序列号
            disk_serial = self.get_disk_serial()
            if disk_serial in ["UNKNOWN", "ERROR"]:
                print("Failed to get disk serial number")
                return False

            # 连接数据库
            conn = mysql.connector.connect(**self.db_config)
            cursor = conn.cursor()

            try:
                # 检查节点是否存在
                check_sql = "SELECT id FROM nodes WHERE nodename = %s"
                cursor.execute(check_sql, (disk_serial,))
                result = cursor.fetchone()

                current_time = datetime.now()
                
                if result:
                    # 更新现有节点
                    update_sql = """
                        UPDATE nodes 
                        SET last_login = %s 
                        WHERE disk_serial = %s
                    """
                    cursor.execute(update_sql, (current_time, disk_serial))
                    print(f"Node {disk_serial} updated")
                else:
                    # 插入新节点
                    insert_sql = """
                        INSERT INTO nodes 
                        (disk_serial, first_login, last_login) 
                        VALUES (%s, %s, %s)
                    """
                    cursor.execute(insert_sql, (disk_serial, current_time, current_time))
                    print(f"New node {disk_serial} registered")

                conn.commit()
                return True

            except Exception as e:
                print(f"Database error: {e}")
                conn.rollback()
                return False
            
            finally:
                cursor.close()
                conn.close()

        except Exception as e:
            print(f"Error in update_node_status: {e}")
            return False

# 使用示例
if __name__ == "__main__":
    # 数据库配置
    db_config = {
        "host": "localhost",
        "user": "your_username",
        "password": "your_password",
        "database": "your_database"
    }

    # 创建节点管理器实例
    node_manager = NodeManager(db_config)
    
    # 更新节点状态
    success = node_manager.update_node_status()
    
    if success:
        print("Node status updated successfully")
    else:
        print("Failed to update node status")